#include "buildingextractor.h"
#include <algorithm>
#include <numeric>
#include "qstring.h"
#include "qdebug.h"
#include <onnxruntime_cxx_api.h>

BuildingExtractor::BuildingExtractor() {}

void BuildingExtractor::initialize(const std::string& modelPath)
{
    env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "BuildingExtractor");
    sessionOptions = Ort::SessionOptions();
    sessionOptions.SetIntraOpNumThreads(15);
    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
    // 是否使用GPU
    //OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);
#ifdef _WIN32
   
    session = Ort::Session(env, QString::fromStdString(modelPath).toStdWString().c_str(), sessionOptions);
#else
    session = Ort::Session(env, modelPath.c_str(), sessionOptions);
#endif

    Ort::AllocatorWithDefaultOptions allocator;
    inputNames.push_back(session.GetInputName(0, allocator));
    outputNames.push_back(session.GetOutputName(0, allocator));

    this->InputShape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    this->OutputShape = session.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
}

cv::Mat BuildingExtractor::preprocess(const cv::Mat& image)
{
    cv::Mat resized, float_image;
    cv::resize(image, resized, cv::Size(1024, 1024));
    resized.convertTo(float_image, CV_32F, 1.0 / 255.0);
    //resized.convertTo(float_image, CV_32F, 1.0 / 65535.0);
    cv::Mat channels[3];
    cv::split(float_image, channels);
   
    for (int i = 0; i < 3; i++) {
        channels[i] = (channels[i] - 0.485) / 0.229;  
    }

    cv::Mat normalized;
    cv::merge(channels, 3, normalized);

    return normalized;
}
std::vector<float> BuildingExtractor::softmax(const std::vector<float>& input)
{
    std::vector<float> result(input.size());
    float max_val = *std::max_element(input.begin(), input.end());
    float sum = 0.0f;

    for (size_t i = 0; i < input.size(); ++i) {
        result[i] = std::exp(input[i] - max_val);
        sum += result[i];
    }

    for (size_t i = 0; i < result.size(); ++i) {
        result[i] /= sum;
    }

    return result;
}

cv::Mat BuildingExtractor::labelToRGB(const cv::Mat& mask)
{
    cv::Mat rgb = cv::Mat::zeros(mask.size(), CV_8UC3);
    rgb.setTo(cv::Scalar(255, 255, 255), mask == 0);
    rgb.setTo(cv::Scalar(0, 0, 0), mask == 1);
    return rgb;
}

cv::Mat BuildingExtractor::extract(const cv::Mat& image)
{
    cv::Mat preprocessed = preprocess(image);

    const int inputHeight = 1024;
    const int inputWidth = 1024;
    const int channels = 3;

    std::vector<int64_t> inputShape{ 1, channels, inputHeight, inputWidth };
    std::vector<float> inputTensorValues(channels * inputHeight * inputWidth);

    // 将预处理后的图像数据复制到输入张量
    for (int c = 0; c < channels; c++) {
        for (int h = 0; h < inputHeight; h++) {
            for (int w = 0; w < inputWidth; w++) {
                inputTensorValues[c * inputHeight * inputWidth + h * inputWidth + w] =
                    preprocessed.at<cv::Vec3f>(h, w)[c];
            }
        }
    }

    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
    auto inputTensor = Ort::Value::CreateTensor<float>(memoryInfo, inputTensorValues.data(), inputTensorValues.size(), inputShape.data(), inputShape.size());

    auto outputTensors = session.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &inputTensor, 1, outputNames.data(), 1);

    const float* mask_data = outputTensors[0].GetTensorMutableData<float>();
    auto outShape = outputTensors[0].GetTensorTypeAndShapeInfo().GetShape();
    int num_cn = outShape[1];
    int out_h = outShape[2];
    int out_w = outShape[3];

    int step = out_h * out_w;
    cv::Mat result = cv::Mat::zeros(cv::Size(out_w, out_h), CV_8UC1);

    for (int row = 0; row < out_h; row++) {
        for (int col = 0; col < out_w; col++) {
            float f1 = mask_data[row * out_w + col];
            float f2 = mask_data[step + row * out_w + col];

            result.at<uchar>(row, col) = (f1 > f2) ? 255 : 0;
        }
    }

    cv::Mat mask;
    cv::resize(result, mask, cv::Size(image.cols, image.rows), 0, 0, cv::INTER_NEAREST);


    cv::Mat output = image.clone();
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    cv::drawContours(output, contours, -1, cv::Scalar(0, 0, 0), 2); 

    return mask;
}